import os
import PyPDF2
import re
import json
import sys
import pyocr
import pyocr.builders
from PIL import Image
import io


# 增加 ocr 识别图片上的内容
tools = pyocr.get_available_tools()
if len(tools) == 0:
    print("No OCR tools found")
    exit(1)

ocr_tool = tools[0]

def convert_pdf_to_text(file_path:str):
    if file_path:
        with open(file_path, 'rb') as pdf_file:
            pdf_reader = PyPDF2.PdfReader(pdf_file)
            num_pages = len(pdf_reader.pages)
            text = ''
            for page_num in range(num_pages):
                page = pdf_reader.pages[page_num]
                if page.extract_text():
                    text += page.extract_text() + " "
                    for img in page.images:
                        image_file = io.BytesIO(img.data)
                        image = Image.open(image_file)
                        img_str = ocr_tool.image_to_string(
                            image,
                            lang='eng',
                            builder=pyocr.builders.TextBuilder()
                        )
                        text += img_str + " "

            text = re.sub(r'\s+', ' ', text).strip()
            sentences = re.split(r'(?<=[.!?]) +', text)
            chunks = []
            current_chunk = ""
            for sentence in sentences:
                if len(current_chunk) + len(sentence) + 1 < 1000:
                    current_chunk += (sentence + " ").strip()
                else:
                    chunks.append(current_chunk)
                    current_chunk = sentence + " "
            if current_chunk:
                chunks.append(current_chunk)
            with open("vault.txt", "a", encoding="utf-8") as vault_file:
                for chunk in chunks:

                    vault_file.write(chunk.strip() + "\n")


def main():
    args = sys.argv
    argc = len(args)
    if argc < 3:
        print("python upload_data.py -pdf [your file path]")
        exit(0)


    if args[1] == "-pdf":
        convert_pdf_to_text(args[2])




if __name__ == "__main__":
    main()
